import torch
import torch.nn as nn
from torchinfo import summary


VGG_Layers = {
    "VGG19": [
        "64", "64", "max_pool",
        "128", "128", "max_pool",
        "256", "256", "256", "256", "max_pool",
        "512", "512", "512", "512", "max_pool",
        "512", "512", "512", "512", "max_pool",
    ],
}

class VGGNetUniverse(nn.Module):
    def __init__(self, vgg_type = "VGG19", in_channels=3, num_classes=100):
        super(VGGNetUniverse, self).__init__()
        # VGG expected 224x224 figure
        size = 224

        layers = []
        features = in_channels
        for layer in VGG_Layers[vgg_type]:
            if layer != "max_pool":
                layers.append(nn.Conv2d(in_channels=features, out_channels=int(layer), kernel_size=3, padding=1, bias=False))
                features = int(layer)
                layers.append(nn.BatchNorm2d(features))
                layers.append(nn.ReLU(inplace=True))
            else:
                layers.append(nn.MaxPool2d(kernel_size=2, stride=2))
                size //= 2

        self.layers = nn.Sequential(*layers)
        self.fc1 = nn.Linear(features * size * size, 4096)
        self.fc2 = nn.Linear(4096, 4096)
        self.fc3 = nn.Linear(4096, num_classes)
        self.dropout = nn.Dropout(0.5)
        self.sigma = nn.ReLU(inplace=True)

    def forward(self, x):
        out = self.layers(x)
        out = out.view(out.size(0), -1)
        out = self.dropout(self.sigma(self.fc1(out)))
        out = self.dropout(self.sigma(self.fc2(out)))
        return self.fc3(out)

def get_model():
    return VGGNetUniverse()

if __name__ == "__main__":
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    model = VGGNetUniverse().to(device=device)
    summary(model, input_size=(256, 3, 224, 224))







